import logging
import time
from typing import Any
import os

from kubernetes import client, config, stream
from kubernetes.client.models import V1Pod
import paramiko
import yaml
from kubernetes.client.rest import ApiException

from common.env import load_copy_task


class StressConfig:
    def __init__(self, cpu_core=0, cpu_load=0, vm=0, vm_bytes=0):
        self.cpu_core = cpu_core
        self.cpu_load = cpu_load
        self.vm = vm
        self.vm_bytes = vm_bytes


class KubeClient:
    def __init__(self, ssh_client: paramiko.SSHClient):
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)
        self.ssh_client = ssh_client
        self.kube_config = self._fetch_and_load_kubeconfig()
        
        _, active_context = config.list_kube_config_contexts(self.kube_config)
        cluster_name = active_context.get('name', '未知')
        self.logger.info(f"当前连接到 K8s 集群: {cluster_name}")

        self.core_v1 = client.CoreV1Api()
        self.apps_v1 = client.AppsV1Api()
        self.custom_api = client.CustomObjectsApi()
        self.static_path = os.path.join(os.path.dirname(__file__), "../static")

    def _fetch_and_load_kubeconfig(self) -> str:
        remote_path, local_path = load_copy_task()
        if os.path.exists(local_path):
            config.load_kube_config(local_path)
            return local_path
        with self.ssh_client.open_sftp() as sftp:
            sftp.get(remote_path, local_path)
        config.load_kube_config(local_path)
        return local_path
        

    def _deep_update(self, d, u):
        for k, v in u.items():
            if isinstance(v, dict) and isinstance(d.get(k), dict):
                self._deep_update(d[k], v)
            else:
                d[k] = v
        return d

    def _install_pod_from_yaml(self, manifest_file: str, name: str = "", request: str = "", limit: str = "") -> V1Pod:
        # 读取yaml文件
        with open(manifest_file) as f:
            pod_manifest = yaml.safe_load(f)
        if name:
            pod_manifest["metadata"]["name"] = name
        if not pod_manifest["metadata"].get("namespace", None):
            pod_manifest["metadata"]["namespace"] = "default"
            
        # 根据是否有request和limit，更新pod_manifest中的值
        # pod 模板保证有这些字段，所以跳过了空值检测
        if request:
            pod_manifest['spec']['containers'][0]['resources']['requests']['cpu'] = request
        if limit:
            pod_manifest['spec']['containers'][0]['resources']['limits']['cpu'] = limit
            
        try:
            pod = self.core_v1.create_namespaced_pod(pod_manifest["metadata"]["namespace"], pod_manifest)
            return pod
        except Exception as e:
            raise RuntimeError(f"Pod创建失败: {str(e)}") from e

    def install_be_pod(self, request: str = "", limit: str = "", name: str = "") -> V1Pod:
        """根据 static/be-pod.yaml 部署 BE Pod, 可以声明 request limit 以覆盖文件中配置，默认为0不覆盖"""
        yaml_path = os.path.join(self.static_path, "be-pod.yaml")
        return self._install_pod_from_yaml(yaml_path, name, request, limit)

    def install_ls_pod(self, request: str = "", limit: str = "", name: str = "") -> V1Pod:
        """根据 static/ls-pod.yaml 部署 LS Pod, 可以声明 request limit 以覆盖文件中配置，默认为0不覆盖"""
        yaml_path = os.path.join(self.static_path, "ls-pod.yaml")
        return self._install_pod_from_yaml(yaml_path, name, request, limit)

    def install_hls_pod(self, req_lim: str = "", name: str = "") -> V1Pod:
        """根据 static/hls-pod.yaml 部署 HLS Pod, 可以声明 request limit 以覆盖文件中配置，默认为0不覆盖"""
        yaml_path = os.path.join(self.static_path, "hls-pod.yaml")
        return self._install_pod_from_yaml(yaml_path, name, req_lim, req_lim)

    def install_common_pod(self, name: str = "") -> V1Pod:
        """只是根据 static/common-pod.yaml 部署一个普通 Pod"""
        yaml_path = os.path.join(self.static_path, "common-pod.yaml")
        return self._install_pod_from_yaml(yaml_path, name)

    def is_pod_existing(self, pod_name, namespace) -> bool:
        try:
            self.core_v1.read_namespaced_pod(pod_name, namespace)
            return True
        except ApiException as e:
            if e.status == 404:
                return False
            else:
                raise e

    def delete_pod(self, pod_name: str, namespace: str, timeout: int = 30):
        if not self.is_pod_existing(pod_name, namespace):
            self.logger.warning(f"Pod {pod_name} 未找到，视为已删除")
            return

        self.core_v1.delete_namespaced_pod(pod_name, namespace)

        for _ in range(timeout):
            if not self.is_pod_existing(pod_name, namespace):
                self.logger.info(f"Pod {pod_name} 已删除")
                return
            time.sleep(1)
        raise TimeoutError(f"Pod {pod_name} 删除超时")
    
    def wait_pod_condition(self, name: str, namespace: str, condition: str, timeout: int = 30) -> bool:
        for _ in range(timeout):
            pod = self.core_v1.read_namespaced_pod(name, namespace)
            if pod.status.phase == condition:
                return True
            time.sleep(1)
        return False

    def wait_pod_delete(self, name: str, namespace: str, timeout: int = 30) -> bool:
        for _ in range(timeout):
            if not self.is_pod_existing(name, namespace):
                return True
            time.sleep(1)
        return False

    def wait_for_volcano_restart(self, timeout: int = 30):
        volcano_label_selector = "app=volcano-scheduler"
        for _ in range(timeout):
            pod_list = self.core_v1.list_namespaced_pod("volcano-system", label_selector=volcano_label_selector)
            if not pod_list.items:
                time.sleep(1)
                continue
            volcano_pod = pod_list.items[0]
            if volcano_pod.status.phase == "Running":
                return True
            time.sleep(1)
        return False

    def list_priorityclass(self) -> dict:
        return self.custom_api.list_cluster_custom_object(
            group="scheduling.k8s.io",
            version="v1",
            plural="priorityclasses",
        )
    
    def deploy_stress_pod(self, hls_resources: Any = None) -> V1Pod:
        """
        安装 stress-ng Pod 并启动 stress-ng 进程。
        默认为 BE Pod，设置HLS 为 True 时以 HLS Pod启动，以避免部分用例中被 rubik 水位线驱逐
        """
        yaml_path = os.path.join(self.static_path, "stress-ng.yaml")
        with open(yaml_path) as f:
            pod_manifest = yaml.safe_load(f)
        if isinstance(hls_resources, dict):
            pod_manifest["metadata"]["labels"].pop("volcano.sh/preemptable")
            pod_manifest["metadata"]["annotations"].pop("volcano.sh/preemptable")
            pod_manifest["metadata"]["annotations"]["openfuyao.com/qos-level"] = "HLS"
            pod_manifest["spec"]["containers"][0]["resources"] = hls_resources
        
        stress_pod = {}
        try:
            stress_pod = self.core_v1.create_namespaced_pod(
                namespace=pod_manifest["metadata"]["namespace"],
                body=pod_manifest,
            )
        except Exception as e:
            raise RuntimeError(f"安装stress-ng Pod失败: {e}") from e

        name, ns = stress_pod.metadata.name, stress_pod.metadata.namespace
        if not self.wait_pod_condition(name, ns, "Running"):
            self.delete_pod(name, ns)
            raise RuntimeError(f"stress-ng 未成功启动，删除中...")

        return stress_pod

    def start_stress_process(self, name: str, namespace: str, stress_config: StressConfig):
        cpu_core = stress_config.cpu_core
        cpu_load = stress_config.cpu_load
        vm = stress_config.vm
        vm_bytes = stress_config.vm_bytes
        stress_pod = {}
        try:
            stress_pod = self.core_v1.read_namespaced_pod(name, namespace)
        except Exception as e:
            raise RuntimeError(f"获取stress-ng Pod失败: {e}") from e

        container = stress_pod.spec.containers[0].name
        command = ["stress-ng"]
        if cpu_core and cpu_load:
            command += [
                "--cpu", str(cpu_core),
                "--cpu-method", "pi",
                "--cpu-load", str(cpu_load),
            ]
        if vm and vm_bytes:
            command += [
                '--vm', str(vm),
                '--vm-bytes', str(vm_bytes) + '%',
            ]
        
        command = " ".join(command) + " > /dev/null 2>&1 &"

        stream.stream(
            self.core_v1.connect_get_namespaced_pod_exec,
            name=stress_pod.metadata.name,
            namespace=stress_pod.metadata.namespace,
            container=container,
            command=["sh", "-c", command],
            stderr=True,
            stdin=False,
            stdout=False,
            tty=False,
        )

        return stress_pod
    
    def stop_stress_process(self, timeout: int = 30):
        # 使用ssh_client查找并关闭stress-ng进程，然后删除Pod
        try:
            # 查找stress-ng进程
            find_cmd = "pgrep stress-ng"
            stdin, stdout, stderr = self.ssh_client.exec_command(find_cmd)
            pid = stdout.read().decode('utf-8').strip().split('\n')[0]
            if pid.isdigit():
                kill_cmd = f"kill -9 {pid}"
                self.ssh_client.exec_command(kill_cmd)
        except Exception as e:
            self.logger.error(f"关闭stress-ng进程时发生异常: {e}")

        for _ in range(timeout):
            find_cmd = "pgrep stress-ng"
            stdin, stdout, stderr = self.ssh_client.exec_command(find_cmd)
            if not stdout.read().decode('utf-8').strip().split('\n'):
                return True
            time.sleep(1)
        return False
